import pod5
import pysam
import numpy as np

pod5_dr=pod5.DatasetReader('/public/home/hpc234701005/data/R9.4/HG002/pod5/output.pod5', recursive=True)
bam_file=pysam.AlignmentFile('/public/home/hpc234701005/data/R9.4/HG002/bam/dorado/mod_infer.sorted.dorado.5mc.full.bam','rb',check_sq=False)
example_result='/public/home/hpc234701005/data/R9.4/HG002/bam/dorado/example_data.tsv'
outputfh = open(example_result, "w", buffering=512)
bam_index=pysam.IndexedReads(bam_file)
bam_index.build()

read_name='4d74fbcb-6c58-4bc5-9918-27ebd8820699'

pod5_record=pod5_dr.get_read(read_name)
signal=pod5_record.signal
read_iter=bam_index.find(read_name)
for bam_read in read_iter:
    tags = dict(bam_read.tags)
    mv_tag = tags["mv"]
    num_trimmed = tags["ts"]
    norm_shift = tags["sm"]
    norm_scale = tags["sd"]
    stride=mv_tag[0]
    mv_table=mv_tag[1:]
    seq=bam_read.query_sequence
    if num_trimmed >= 0:
        signal_trimmed = (
            signal[num_trimmed:] - norm_shift
        ) / norm_scale
    else:
        signal_trimmed = (
            signal[:num_trimmed] - norm_shift
        ) / norm_scale
    sshift, sscale = np.mean(signal_trimmed), float(np.std(signal_trimmed))
    if sscale == 0.0:
        norm_signals = signal_trimmed
    else:
        norm_signals = (signal_trimmed - sshift) / sscale
    mm=tags["MM"]
    ml=tags["ML"]

    outputfh.write(
                    f"{seq}\t{np.array2string(norm_signals, threshold=1000)}\t"
                    f"{mv_tag}\t{mm}\t{ml}\n"
                    )
    break
    